FROM {{ cuda.base_image }} as base

{% if "x86_64" in cuda %}
FROM base as base-amd64

ENV NVARCH x86_64

{% if "requires" in cuda.x86_64 %}
ENV NVIDIA_REQUIRE_CUDA "{{ cuda.x86_64.requires }}"
{% endif %}
ENV NV_CUDA_CUDART_VERSION {{ cuda.x86_64.components.cudart.version }}

{% endif %}
{% if "arm64" in cuda %}
FROM base as base-arm64

{% if "flavor" in cuda and cuda.flavor == "jetson" %}
ENV NVARCH {{ cuda.supported_arches.tegra.repo_name.deb }}
{% else %}
ENV NVARCH sbsa
{% endif %}
{% if "requires" in cuda.arm64 %}
ENV NVIDIA_REQUIRE_CUDA "{{ cuda.arm64.requires }}"
{% endif %}
ENV NV_CUDA_CUDART_VERSION {{ cuda.arm64.components.cudart.version }}

{% endif %}
{% if "ppc64le" in cuda %}
FROM base as base-ppc64le

ENV NVARCH ppc64el
{% if "requires" in cuda.ppc64le %}
ENV NVIDIA_REQUIRE_CUDA "{{ cuda.ppc64le.requires }}"
{% endif %}
ENV NV_CUDA_CUDART_VERSION {{ cuda.ppc64le.components.cudart.version }}

{% endif -%}

FROM base-${TARGETARCH}

ARG TARGETARCH

LABEL maintainer "NVIDIA CORPORATION <cudatools@nvidia.com>"

{% if "flavor" in cuda and cuda.flavor == "jetson" %}
ARG NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS=base-only
ENV NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS ${NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS}

{% endif -%}

RUN apt-get update &&{% if cuda.flavor == "jetson" %} apt upgrade -y &&{% endif %} apt-get install -y --no-install-recommends \
{% if cuda.os.version == "22.04" and not (cuda.image_tag_suffix | length) %}
    gnupg2 curl ca-certificates && \
    curl -fsSLO {{ cuda.repo_url }}/${NVARCH}/cuda-keyring_1.1-1_all.deb && \
    dpkg -i cuda-keyring_1.1-1_all.deb && \
    apt-get purge --autoremove -y curl \
{% elif cuda.os.version != "16.04" or cuda.flavor == "jetson" %}
    {# Still used for internal 22.04 images until https://jirasw.nvidia.com/browse/CUDAINST-1240 #}
    gnupg2 curl ca-certificates && \
    curl -fsSL {{ cuda.repo_url }}/${NVARCH}/3bf863cc.pub | apt-key add - && \
    echo "deb {{ cuda.repo_url }}/${NVARCH} /" > /etc/apt/sources.list.d/cuda.list && \
    apt-get purge --autoremove -y curl \
{% else %}
    ca-certificates apt-transport-https gnupg-curl && \
    NVIDIA_GPGKEY_SUM=a21c1a0b18a4196fa901b833e13c4fa64f094d7d9e8a6495318e7255f0ef23d1 && \
    NVIDIA_GPGKEY_FPR=eb693b3035cd5710e231e123a4b469963bf863cc && \
    apt-key adv --fetch-keys {{ cuda.repo_url }}/${NVARCH}/3bf863cc.pub && \
    apt-key adv --export --no-emit-version -a $NVIDIA_GPGKEY_FPR | tail -n +5 > cudasign.pub && \
    echo "$NVIDIA_GPGKEY_SUM  cudasign.pub" | sha256sum -c --strict - && rm cudasign.pub && \
    echo "deb {{ cuda.repo_url }}/${NVARCH} /" > /etc/apt/sources.list.d/cuda.list && \
    apt-get purge --auto-remove -y gnupg-curl \
{% endif %}
    && rm -rf /var/lib/apt/lists/*

ENV CUDA_VERSION {{ cuda.version.release_label }}

# For libraries in the cuda-compat-* package: https://docs.nvidia.com/cuda/eula/index.html#attachment-a
RUN apt-get update && apt-get install -y --no-install-recommends \
    cuda-cudart-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_CUDART_VERSION} \
    cuda-compat-{{ cuda.version.major }}-{{ cuda.version.minor }} \
{% if ( cuda.version.major | int ) == 11 and ( cuda.version.minor | int ) <= 2 %}
    && ln -s cuda-{{ cuda.version.major }}.{{ cuda.version.minor }} /usr/local/cuda \
{% endif %}
    && rm -rf /var/lib/apt/lists/*

# Required for nvidia-docker v1
RUN echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf \
    && echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf

ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH}
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64

COPY NGC-DL-CONTAINER-LICENSE /

# nvidia-container-runtime
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
